library(tidyr)
library(dplyr)
library(stringr)
library(ggplot2)

if (!file.exists("output")) {
  dir.create("output")
}

# load data
source("HNC_metadata_tidy.R")
source("data_handling.R")

cvCOSMIC <- read.csv("Input_signature_attributions/Input_CN_COSMIC_attributions.csv", stringsAsFactors = T) %>% rename(donor_id = X)

additional_df <- list(cvCOSMIC,
  # CN clusters generated in Figure_7c.R
  CN_clusters = read.csv("output/CN_clusters.csv"),
  # cancer genes with driver mutations or CN alterations per sample
  drivers_cnm = read.csv("ExtendedDataFigure_10d_data_CN_cancergenes.csv") %>%
    select("donor_id", "MYC", "FADD") %>% rename_with(~ paste0(., "_gain"), .cols = c("MYC", "FADD")) %>% mutate_if(is.numeric, factor),
  # cancer genes with driver mutations per sample. Selecting only genes with >4 positive samples
  drivers_m = read.csv("ExtendedDataFigure_10d_data_mutated_cancergenes.csv") %>%
    select(names(.)[c(T, colSums(.[, -1]) >= 4)]) %>% mutate_if(is.numeric, factor)
)

for (df in additional_df) {
  data <- data %>% left_join(df, by = "donor_id")
}

data <- data %>% filter(!is.na(HNC_clusters))

# Select signatures present in >n% of samples
# Dichotomize the signatures (Presence (signatures>0)="1",absence (signatures==0)="0")
# If signature is present in more than 75% cases, dichotomize by the median
data <- dichotomizeSigs(metadata = data, sigsn = cvCOSMIC, sigcutoff = 2)

## dummify cluster variables
data <- data %>%
  mutate(clusters_dummy = 1) %>%
  spread(HNC_clusters, clusters_dummy) %>%
  mutate_at(vars(levels(factor(data$HNC_clusters))), \(x){
    replace_na(x, 0)
  }) %>%
  mutate(clusters_dummy = 1) %>%
  spread(k2, clusters_dummy) %>%
  mutate_at(vars(levels(factor(data$k2))), \(x){
    replace_na(x, 0)
  }) %>%
  select(-c("Cluster D", "Cluster P1", "Cluster P2"))

signatures <- c(grep("Cluster", names(data), value = T), grep("^CN.*_cat$", names(data), value = T))
factors <- c(names(additional_df$drivers_m)[-1], names(additional_df$drivers_cnm)[-1])

fisher_results <- NULL
for (f in factors) {
  for (s in signatures) {
    if (exists("test") == T) {
      rm(test)
    }
    t <- table(data[, s], data[, f])
    if (ncol(t) < 2 | nrow(t) < 2) {
      print(paste("Variable", f, "for signature", s, "has been removed due to <2 categories"))
    } else {
      test <- tryCatch(
        {
          fisher.test(t)
        },
        error = function(e) {
          NULL
        }
      )

      if (is.null(test)) {
        print(paste0("Variable ", f, ", signature ", s, ": removed due to ERROR in fisher test"))
        next
      }

      result <- data.frame(
        signature = s,
        factor = f,
        category1 = levels(factor(data[, f]))[2],
        category2 = levels(factor(data[, f]))[1],
        perc_in_factor_c1 = t[2, 2] / (t[2, 2] + t[1, 2]) * 100,
        prop_in_factor_c1 = paste0(t[2, 2], "/", t[2, 2] + t[1, 2]),
        perc_in_factor_c2 = t[2, 1] / (t[2, 1] + t[1, 1]) * 100,
        prop_in_factor_c2 = paste0(t[2, 1], "/", t[2, 1] + t[1, 1]),
        OR = (t[2, 2] / t[1, 2]) / (t[2, 1] / t[1 / 1]),
        fisher_p = test$p.value
      )
      fisher_results <- rbind(fisher_results, result)
    }
  }
}

# tidy fisher test results for plotting
fisher.plot <- fisher_results %>%
  group_by(factor) %>%
  mutate(
    signature = factor(signature, levels = signatures),
    q_val = p.adjust(fisher_p, "fdr"),
    q_val_log2 = -log(q_val, 2),
    log_OR = log(OR)
  ) %>%
  ungroup() %>%
  select(signature, factor, OR, log_OR, fisher_p, q_val, q_val_log2) %>%
  filter(q_val < 0.05)

order_x <- c("PTEN", "RB1", "CASP8", "HRAS", "SETD2", "TP53", "FADD_gain", "MYC_gain")
order_y <- c(signatures[signatures %in% unique(fisher.plot$signature)])

ggplot(fisher.plot, aes(x = factor, y = signature, color = log_OR, size = q_val_log2)) +
  geom_point() +
  labs(y = "CNV signatures", x = "Driver genes", color = "log(OR)", size = "-log2(q)") +
  scale_x_discrete(limits = order_x, labels = str_replace_all(order_x, "_", " ")) +
  scale_y_discrete(limits = rev(order_y), labels = str_remove(rev(order_y), "_cat")) +
  scale_color_gradient2(high = "#ca0020", mid = "white", low = "#0571b0", midpoint = 0, oob = scales::squish_infinite) +
  scale_size_continuous(limits = c(min(fisher.plot$q_val_log2), max(fisher.plot$q_val_log2))) +
  theme_bw() +
  theme(
    axis.text.x = element_text(angle = 45, hjust = 1),
    panel.grid = element_blank()
    )

ggsave(last_plot(),
  filename = paste0("output/ExtendedDataFigure_10d_", Sys.Date(), ".pdf"),
  device = "pdf", width = 4, height = length(unique(fisher.plot$signature)) / 3, units = "in"
)
